import os
import asyncio
from typing import List, Dict, Any, Optional
from datetime import datetime
from pathlib import Path
import markdown
import json

from ..models.kimi_client import KimiClient
from ..tools.document_processor import DocumentProcessor, DocumentChunk

class ReportGenerator:
    """
    研究报告生成器
    """
    
    def __init__(self, kimi_client: Optional[KimiClient] = None):
        self.kimi_client = kimi_client or KimiClient()
        self.processor = DocumentProcessor()
        self.report_config = {
            "language": os.getenv("REPORT_LANGUAGE", "zh"),
            "format": os.getenv("REPORT_FORMAT", "markdown"),
            "max_length": int(os.getenv("MAX_REPORT_LENGTH", 5000))
        }
    
    async def generate_research_report(
        self,
        topic: str,
        search_results: List[Dict[str, Any]],
        report_type: str = "comprehensive",
        custom_sections: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        """
        生成研究报告
        
        Args:
            topic: 研究主题
            search_results: 搜索结果列表
            report_type: 报告类型 ("comprehensive", "brief", "technical")
            custom_sections: 自定义章节列表
            
        Returns:
            包含报告内容的字典
        """
        print(f"开始生成关于{topic}的研究报告...")
        
        # 处理搜索内容
        all_content = []
        for result in search_results:
            content = result.get("content", "")
            if content:
                chunks = self.processor.chunk_text(content, result.get("title", "unknown"))
                all_content.extend(chunks)
        
        # 提取关键词
        keywords = self._extract_keywords(topic)
        
        # 生成报告结构
        report_structure = self._create_report_structure(
            topic, report_type, custom_sections
        )
        
        # 生成各章节内容
        report_sections = {}
        
        for section_name, section_info in report_structure.items():
            print(f"正在生成章节: {section_info['title']}")
            
            context = self._prepare_section_context(
                section_name, all_content, keywords
            )
            
            section_content = await self._generate_section_content(
                topic, section_name, context, section_info
            )
            
            report_sections[section_name] = {
                "title": section_info["title"],
                "content": section_content,
                "order": section_info["order"]
            }
        
        # 生成完整的报告
        full_report = self._compile_report(topic, report_sections)
        
        # 生成元数据
        metadata = {
            "topic": topic,
            "report_type": report_type,
            "generated_at": datetime.now().isoformat(),
            "sources_count": len(search_results),
            "total_word_count": len(full_report),
            "sections_count": len(report_sections)
        }
        
        return {
            "content": full_report,
            "sections": report_sections,
            "metadata": metadata,
            "sources": search_results
        }
    
    def _extract_keywords(self, topic: str) -> List[str]:
        """
        从主题中提取关键词
        
        Args:
            topic: 研究主题
            
        Returns:
            关键词列表
        """
        # 简单的关键词提取
        keywords = [topic]
        
        # 添加常见相关词
        common_words = ["研究", "分析", "发展", "现状", "趋势", "应用", "技术"]
        for word in common_words:
            if word not in topic:
                keywords.append(f"{topic}{word}")
                keywords.append(f"{word}{topic}")
        
        # 分割主题词
        if len(topic) > 4:
            parts = topic.split()
            keywords.extend(parts)
        
        return list(set(keywords))
    
    def _create_report_structure(
        self,
        topic: str,
        report_type: str,
        custom_sections: Optional[List[str]] = None
    ) -> Dict[str, Dict[str, Any]]:
        """
        创建报告结构
        
        Args:
            topic: 研究主题
            report_type: 报告类型
            custom_sections: 自定义章节
            
        Returns:
            报告结构字典
        """
        base_structure = {
            "introduction": {
                "title": "引言",
                "description": "介绍研究主题的背景和重要性",
                "order": 1
            },
            "background": {
                "title": "背景知识",
                "description": "主题相关的基本概念和历史发展",
                "order": 2
            },
            "current_status": {
                "title": "现状分析",
                "description": "当前主题的发展现状和主要特点",
                "order": 3
            },
            "key_findings": {
                "title": "主要发现",
                "description": "研究中的重要发现和结论",
                "order": 4
            },
            "applications": {
                "title": "应用与影响",
                "description": "主题的实际应用和社会影响",
                "order": 5
            },
            "future_trends": {
                "title": "未来趋势",
                "description": "主题未来的发展方向和趋势预测",
                "order": 6
            },
            "conclusion": {
                "title": "总结",
                "description": "研究总结和建议",
                "order": 7
            }
        }
        
        if report_type == "brief":
            # 简版报告
            return {
                "overview": {
                    "title": "概览",
                    "description": "主题概览",
                    "order": 1
                },
                "key_points": {
                    "title": "要点",
                    "description": "主要观点",
                    "order": 2
                },
                "summary": {
                    "title": "总结",
                    "description": "简要总结",
                    "order": 3
                }
            }
        elif report_type == "technical":
            # 技术报告
            tech_structure = base_structure.copy()
            tech_structure.update({
                "methodology": {
                    "title": "研究方法",
                    "description": "采用的研究方法和技术路线",
                    "order": 3
                },
                "data_analysis": {
                    "title": "数据分析",
                    "description": "详细的数据分析和结果",
                    "order": 4
                }
            })
            return tech_structure
        
        if custom_sections:
            # 添加自定义章节
            custom_structure = base_structure.copy()
            for i, section in enumerate(custom_sections, 8):
                custom_structure[f"custom_{i}"] = {
                    "title": section,
                    "description": f"关于{section}的详细分析",
                    "order": i
                }
        
        return base_structure
    
    def _prepare_section_context(
        self,
        section_name: str,
        chunks: List[DocumentChunk],
        keywords: List[str]
    ) -> str:
        """
        为章节准备上下文
        
        Args:
            section_name: 章节名称
            chunks: 文档块列表
            keywords: 关键词列表
            
        Returns:
            上下文文本
        """
        # 根据章节类型选择相关内容
        section_keywords = {
            "introduction": ["介绍", "背景", "重要性", "意义"],
            "background": ["历史", "发展", "概念", "定义"],
            "current_status": ["现状", "目前", "当前", "发展"],
            "key_findings": ["发现", "结论", "结果", "研究表明"],
            "applications": ["应用", "使用", "实践", "案例"],
            "future_trends": ["未来", "趋势", "预测", "展望"],
            "conclusion": ["总结", "建议", "展望", "结论"]
        }
        
        section_keyword_list = section_keywords.get(section_name, keywords)
        
        # 选择与章节相关的文档块
        relevant_chunks = self.processor.organize_by_relevance(
            chunks, section_keyword_list
        )
        
        # 选择最相关的几个块
        selected_chunks = relevant_chunks[:5]
        
        # 合并内容
        context_parts = [chunk.content for chunk in selected_chunks]
        return " ".join(context_parts)
    
    async def _generate_section_content(
        self,
        topic: str,
        section_name: str,
        context: str,
        section_info: Dict[str, Any]
    ) -> str:
        """
        生成章节内容
        
        Args:
            topic: 研究主题
            section_name: 章节名称
            context: 上下文内容
            section_info: 章节信息
            
        Returns:
            章节内容
        """
        section_prompts = {
            "introduction": f"请为{topic}主题的研究报告写一个引人入胜的引言。介绍该主题的重要性、研究背景和研究目的。基于以下信息：{context}",
            "background": f"请提供{topic}主题的详细背景知识。包括基本概念、历史发展、相关理论和重要里程碑。基于以下信息：{context}",
            "current_status": f"请分析{topic}主题的当前发展现状。包括主要特点、市场规模、技术水平和竞争格局。基于以下信息：{context}",
            "key_findings": f"请总结关于{topic}的主要研究发现和重要结论。基于以下信息：{context}",
            "applications": f"请详细说明{topic}的实际应用场景、成功案例和社会经济影响。基于以下信息：{context}",
            "future_trends": f"请预测{topic}主题的未来发展趋势。包括技术发展方向、市场机遇和潜在挑战。基于以下信息：{context}",
            "conclusion": f"请为{topic}主题的研究报告提供总结。概括主要观点，并提出建设性建议。基于以下信息：{context}"
        }
        
        prompt = section_prompts.get(section_name, 
            f"请为{topic}主题生成关于{section_info['title']}的内容。基于以下信息：{context}")
        
        messages = [{"role": "user", "content": prompt}]
        
        try:
            content = await self.kimi_client.chat_completions(
                messages, temperature=0.7, max_tokens=1000
            )
            return content
        except Exception as e:
            print(f"生成章节{section_name}失败: {e}")
            return f"# {section_info['title']}\n\n该章节内容生成失败，请稍后重试。"
    
    def _compile_report(self, topic: str, sections: Dict[str, Dict[str, Any]]) -> str:
        """
        编译完整报告
        
        Args:
            topic: 研究主题
            sections: 章节字典
            
        Returns:
            完整报告文本
        """
        # 按顺序排序章节
        sorted_sections = sorted(
            sections.items(), 
            key=lambda x: x[1]["order"]
        )
        
        # 生成标题
        report_title = f"# {topic} 研究报告\n\n"
        report_date = f"*生成时间: {datetime.now().strftime('%Y年%m月%d日')}*\n\n"
        
        # 生成目录
        toc = "## 目录\n\n"
        for section_name, section_info in sorted_sections:
            toc += f"- [{section_info['title']}](#{section_name})\n"
        toc += "\n"
        
        # 生成内容
        content = ""
        for section_name, section_info in sorted_sections:
            content += f"## {section_info['title']}\n\n"
            content += f"{section_info['content']}\n\n"
        
        return report_title + report_date + toc + content
    
    async def save_report(self, report: Dict[str, Any], filename: Optional[str] = None) -> str:
        """
        保存报告到文件
        
        Args:
            report: 报告字典
            filename: 文件名（可选）
            
        Returns:
            文件路径
        """
        if not filename:
            topic = report["metadata"]["topic"]
            safe_topic = "".join(c for c in topic if c.isalnum() or c in (' ', '-', '_')).rstrip()
            safe_topic = safe_topic.replace(' ', '_')
            filename = f"{safe_topic}_研究报告_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
        
        # 确保目录存在
        output_dir = Path("reports")
        output_dir.mkdir(exist_ok=True)
        
        file_path = output_dir / filename
        
        # 保存Markdown文件
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(report["content"])
        
        # 保存元数据
        metadata_file = output_dir / f"{file_path.stem}_metadata.json"
        with open(metadata_file, 'w', encoding='utf-8') as f:
            json.dump(report["metadata"], f, ensure_ascii=False, indent=2)
        
        # 保存来源信息
        sources_file = output_dir / f"{file_path.stem}_sources.json"
        with open(sources_file, 'w', encoding='utf-8') as f:
            json.dump(report["sources"], f, ensure_ascii=False, indent=2)
        
        print(f"报告已保存到: {file_path}")
        return str(file_path)